from fastapi import FastAPI, Request
from flask import request
import requests
from pydantic import BaseModel
from typing import Optional
import uvicorn
# 初始化 FastAPI 应用
app = FastAPI()


from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import chat_agent_executor

from model import get_model

# sqlalchemy, 初始化 MySql 数据库连接
HOST = 'rm-2vcy212o553nqof7xpo.mysql.cn-chengdu.rds.aliyuncs.com'
PORT = '3306'
DATABASE = 'prod-platform'
USER = 'ygkj'
PASSWORD = 'ki854a#!acx'

MYSQL_URI = 'mysql+pymysql://{user}:{password}@{host}:{port}/{db}?charset=utf8'.format(
    user=USER, password=PASSWORD, host=HOST, port=PORT, db=DATABASE)

# 1. 连接 MySQL 数据库
db = SQLDatabase.from_uri(MYSQL_URI)

# 2. 调用大模型
zhipuai = get_model()

# 创建工具
toolkit = SQLDatabaseToolkit(db=db, llm=zhipuai)
tools = toolkit.get_tools()

# 提示词
system_prompt = """
你是一个被设计用来与SQL数据库交互的代理。
将给定的输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。

你可以使用与数据库交互的工具。在执行查询之前，你必须仔细检查。如果在执行查询时出现错误，请重写查询并重试。
不要对数据库做任何DML语句（插入，更新，删除等）

首先，你应该查看数据库中的表，看看可以查询什么。不要跳过这一步。
然后查询最相关的表的模式。
"""
system_message = SystemMessage(content=system_prompt)

# 创建代理
agent_executor = chat_agent_executor.create_tool_calling_executor(zhipuai, tools, state_modifier=system_message)

@app.post("/query")
async def query_info(request: Request):
    # 从请求体中获取 JSON 数据
    json_data = await request.json()
    #
    # business = json_data.get("business")
    # platform = json_data.get("platform")
    # place = json_data.get("place")
    # req = json_data.get("request")
    #
    # # 构建查询字符串
    # query = f"业务类型: {business}, 平台: {platform}, 地区: {place}, 需求: {req}"

    query = json_data.get("input")  # 提取 input 字段
    print(f'query:{query}')

    if not query:
        return {"error": "Missing 'input' field in the request body"}

    # 调用 agent_executor 并返回结果
    result = agent_executor.invoke({'messages': [HumanMessage(content=query)]})['messages'][-1].content
    print(result)
    return {"response": result}

@app.get("/query")
async def query_info():
    # 从请求体中获取 JSON 数据
    # json_data = await request.json()
    # query = json_data.get("input")  # 提取 input 字段
    # if not query:
    #     return {"error": "Missing 'input' field in the request body"}
    query = "全国云剪"
    # 调用 agent_executor 并返回结果
    result = agent_executor.invoke({'messages': [HumanMessage(content=query)]})['messages'][-1].content
    print(result)
    return {"response": result}
if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=3433)
